import shutil
from pathlib import Path
from typing import Optional

SOURCE_DIR = "/hdd/data/swp/casia/liuchanhg"
DEST_BASE = "/hdd/data/swp/DH/CASIA/audio"


NEGATIVE_DIR_NAMES = {"angry", "fear", "sad"}
POSITIVE_DIR_NAMES = {"happy", "surprise"}
NEUTRAL_DIR_NAMES = {"neutral"}


def ensure_dest_dirs(base_dir: str) -> None:
    for sub in {"negative", "positive", "neutral"}:
        Path(base_dir, sub).mkdir(parents=True, exist_ok=True)


def resolve_category_from_dir(dir_name: str) -> Optional[str]:
    name = dir_name.strip().lower()
    if name in NEGATIVE_DIR_NAMES:
        return "negative"
    if name in POSITIVE_DIR_NAMES:
        return "positive"
    if name in NEUTRAL_DIR_NAMES:
        return "neutral"
    return None


def copy_casia_by_subdir(source_dir: str, dest_base: str) -> None:
    ensure_dest_dirs(dest_base)

    src_path = Path(source_dir)
    if not src_path.exists():
        raise FileNotFoundError(f"Source directory not found: {source_dir}")

    # prefix to add to each wav filename, based on source root directory name
    filename_prefix = src_path.name.strip()

    total = 0
    copied = 0
    skipped = 0
    unknown_dir = 0

    # iterate only immediate subdirectories which indicate the emotion
    for subdir in [p for p in src_path.iterdir() if p.is_dir()]:
        category = resolve_category_from_dir(subdir.name)
        if not category:
            unknown_dir += 1
            print(f"[SKIP-DIR] Unknown emotion dir: {subdir}")
            continue

        dest_dir = Path(dest_base, category)

        for wav in subdir.rglob("*.wav"):
            total += 1
            # filename format: <source_root>_<emotion_subdir>_<original_name>
            new_name = f"{filename_prefix}_{subdir.name}_{wav.name}"
            dest_file = dest_dir / new_name
            try:
                shutil.copy2(wav, dest_file)
                copied += 1
                print(f"[OK] {wav} -> {dest_file}")
            except Exception as e:
                skipped += 1
                print(f"[ERR] Failed to copy {wav} -> {dest_file}: {e}")

    print(
        f"\nDone. Files scanned: {total}, Copied: {copied}, File skips: {skipped}, Unknown subdirs: {unknown_dir}.\n"
        f"Dest base: {dest_base}"
    )


if __name__ == "__main__":
    copy_casia_by_subdir(SOURCE_DIR, DEST_BASE)


